import inspect
import operator
from typing import Any, Callable, Dict, List, Optional, Tuple, Union


class Node:
    def __init__(self, name: str, op: str, target: Any, args: tuple, kwargs: dict):
        self.name = name
        self.op = op
        self.target = target
        self.args = args
        self.kwargs = kwargs
        self.users: List[Node] = []

    def __repr__(self):
        return f"Node({self.name}, op={self.op}, target={self.target})"


class Graph:
    def __init__(self):
        self.nodes: List[Node] = []
        self.node_counter = 0
        self.input_nodes: List[Node] = []
        self.output_node: Optional[Node] = None

    def add_node(self, node: Node) -> Node:
        self.nodes.append(node)

        # 更新用户关系
        for arg in node.args:
            if isinstance(arg, Node):
                arg.users.append(node)

        for kwarg in node.kwargs.values():
            if isinstance(kwarg, Node):
                kwarg.users.append(node)

        # 记录特殊节点
        if node.op == "placeholder":
            self.input_nodes.append(node)
        elif node.op == "output":
            self.output_node = node

        return node

    def create_node(self, op: str, target: Any, args: tuple, kwargs: dict) -> Node:
        name = f"{op}_{self.node_counter}"
        self.node_counter += 1
        node = Node(name, op, target, args, kwargs)
        return self.add_node(node)

    def placeholder(self, name: str) -> Node:
        return self.create_node("placeholder", name, (), {})

    def output(self, result: Any) -> Node:
        return self.create_node("output", "output", (result,), {})

    def to_graphviz(self, title: str = "Computational Graph") -> str:
        dot = [
            "digraph G {",
            "  rankdir=TB;",  # 从上到下布局
            "  node [shape=box];",  # 默认矩形节点
            f'  label="{title}";',
            "",
        ]

        # 添加输入节点
        for node in self.input_nodes:
            dot.append(f'  "{node.name}" [label="{node.target}", shape=circle];')

        # 添加计算节点
        for node in self.nodes:
            if node.op == "placeholder" or node.op == "output":
                continue

            # 简化目标名称显示
            target_name = str(node.target)
            if hasattr(node.target, "__name__"):
                target_name = node.target.__name__
            elif callable(node.target):
                target_name = node.target.__class__.__name__

            # 只显示计算类型
            dot.append(f'  "{node.name}" [label="{target_name}"];')

        # 添加输出节点
        if self.output_node:
            output_name = self.output_node.name

            dot.append(
                f'  "{self.output_node.name}" [label="{output_name}", shape=doublecircle];'
            )

        # 添加边连接关系
        dot.append("")
        dot.append("  /* 拓扑关系 */")

        # 遍历所有节点添加边
        for node in self.nodes:
            if node.op == "placeholder":  # 跳过输入节点
                continue

            # 添加参数边
            for arg in node.args:
                if isinstance(arg, Node):
                    dot.append(f'  "{arg.name}" -> "{node.name}";')

            # 添加关键字参数边
            for v in node.kwargs.values():
                if isinstance(v, Node):
                    dot.append(f'  "{v.name}" -> "{node.name}";')

        # 添加到输出节点的边
        if self.output_node and self.output_node.args:
            output_arg = self.output_node.args[0]
            if isinstance(output_arg, Node):
                dot.append(f'  "{output_arg.name}" -> "{self.output_node.name}";')

        # 添加层级排序
        dot.append("")
        dot.append("  /* 层级排序 */")
        dot.append(
            "  { rank=same; "
            + "; ".join(f'"{n.name}"' for n in self.input_nodes)
            + " }"
        )
        if self.output_node:
            dot.append("  { rank=sink; " + f'"{self.output_node.name}"' + " }")

        dot.append("}")
        return "\n".join(dot)

    def save_graphviz(self, filename: str, title: str = "Computational Graph"):
        """保存Graphviz DOT文件"""
        dot_content = self.to_graphviz(title)
        with open(filename, "w") as f:
            f.write(dot_content)
        print(f"Graph saved to {filename}. Use Graphviz to render it.")


class Proxy:
    def __init__(self, node: Node, graph: Graph, tracer: "Trace"):
        self.node = node
        self.graph = graph
        self.tracer = tracer

    def __call__(self, *args, **kwargs):
        args = self.tracer.unwrap(args)
        kwargs = self.tracer.unwrap(kwargs)
        return self.tracer.create_proxy("call_function", self.node.target, args, kwargs)

    # 运算符重载
    def __add__(self, other):
        return self._binop(other, operator.add)

    def __sub__(self, other):
        return self._binop(other, operator.sub)

    def __mul__(self, other):
        return self._binop(other, operator.mul)

    def __truediv__(self, other):
        return self._binop(other, operator.truediv)

    def __pow__(self, other):
        return self._binop(other, operator.pow)

    def __eq__(self, other):
        return self._binop(other, operator.eq)

    def __lt__(self, other):
        return self._binop(other, operator.lt)

    def __gt__(self, other):
        return self._binop(other, operator.gt)

    def __radd__(self, other):
        return self._binop(other, operator.add, reverse=True)

    def __rsub__(self, other):
        return self._binop(other, operator.sub, reverse=True)

    def __rmul__(self, other):
        return self._binop(other, operator.mul, reverse=True)

    def __rtruediv__(self, other):
        return self._binop(other, operator.truediv, reverse=True)

    def __neg__(self):
        return self._unaryop(operator.neg)

    def __abs__(self):
        return self._unaryop(operator.abs)

    # 矩阵乘法
    def __matmul__(self, other):
        return self._binop(other, operator.matmul)

    def __rmatmul__(self, other):
        return self._binop(other, operator.matmul, reverse=True)

    # 容器操作
    def __getitem__(self, key):
        key = self.tracer.unwrap(key)
        return self.tracer.create_proxy(
            "call_function", operator.getitem, (self.node, key), {}
        )

    # 辅助方法
    def _binop(self, other, op, reverse=False):
        other = self.tracer.unwrap(other)
        if reverse:
            args = (other, self.node)
        else:
            args = (self.node, other)
        return self.tracer.create_proxy("call_function", op, args, {})

    def _unaryop(self, op):
        return self.tracer.create_proxy("call_function", op, (self.node,), {})


class Trace:
    def __init__(self):
        self.graph = Graph()
        self.root = None
        self.proxy_map: Dict[Any, Proxy] = {}

    def unwrap(self, value: Any) -> Any:
        if isinstance(value, Proxy):
            return value.node
        elif isinstance(value, (list, tuple)):
            return type(value)(self.unwrap(v) for v in value)
        elif isinstance(value, dict):
            return {k: self.unwrap(v) for k, v in value.items()}
        return value

    def create_proxy(self, op: str, target: Any, args: tuple, kwargs: dict) -> Proxy:
        node = self.graph.create_node(op, target, args, kwargs)
        proxy = Proxy(node, self.graph, self)
        self.proxy_map[node] = proxy
        return proxy

    def trace(self, root: Callable) -> Graph:
        self.root = root
        signature = inspect.signature(root)

        input_proxies = []
        for param_name in signature.parameters:
            placeholder = self.graph.placeholder(param_name)
            proxy = Proxy(placeholder, self.graph, self)
            input_proxies.append(proxy)

        result = root(*input_proxies)
        self.graph.output(self.unwrap(result))
        return self.graph


def symbolic_trace(fn: Callable) -> Graph:
    tracer = Trace()
    return tracer.trace(fn)


# 示例使用 - 创建计算图
def complex_calculation(x, y, z):
    # 基本运算
    a = x + y
    b = a * z
    c = b - 3.5
    d = c**2
    e = -d

    # 矩阵运算
    f = e @ x

    # 条件判断
    g = f if (x > y) else z

    return g


# 跟踪函数并生成图形
graph = symbolic_trace(complex_calculation)
graph.save_graphviz("computation_graph.dot", title="Computational Topology")

print("\nGraphviz DOT file generated. Use the following command to render:")
print("  dot -Tpng computation_graph.dot -o computation_graph.png")
print("  dot -Tsvg computation_graph.dot -o computation_graph.svg")
